package org.acghub.common.topology;

import lombok.Getter;

import java.util.*;

/**
 * 有向无环图，环路检测器
 *
 * @param <T>
 */
public class DAGCircleFinder<T> implements CircleFinder<T> {

    @Getter
    protected final Map<T, List<T>> adj;

    protected final Map<T, VisitState> visit = new HashMap<>();

    protected final Map<T, T> child2Father = new HashMap<>();

    protected final List<T> circle = new ArrayList<>();

    public DAGCircleFinder(Map<T, List<T>> adj) {
        this.adj = adj;
    }

    public DAGCircleFinder() {
        this.adj = new HashMap<>();
    }

    @Override
    public void addEdge(T from, T to) {
        Objects.requireNonNull(from, "from is null");
        Objects.requireNonNull(to, "to is null");
        adj.computeIfAbsent(from, k -> new ArrayList<>()).add(to);
    }

    @Override
    public boolean removeEdge(T from, T to) {
        Objects.requireNonNull(from, "from is null");
        Objects.requireNonNull(to, "to is null");
        for (Map.Entry<T, List<T>> entry : adj.entrySet()) {
            if (entry.getKey().equals(from)) {
                List<T> toList = entry.getValue();
                if (toList.remove(to)) {
                    return true;
                }
            }
        }
        return false;
    }

    @Override
    public FindCircleResp<T> findCircle() {
        for (T u : adj.keySet()) {
            if (visit.getOrDefault(u, VisitState.WHITE) == VisitState.WHITE) {
                dfs(u);
            }
            if (!circle.isEmpty()) {
                break;
            }
        }
        return new FindCircleResp<>(circle);
    }

    private void dfs(T u) {
        if (!circle.isEmpty()) {
            return;
        }
        visit.put(u, VisitState.GRAY);
        for (T v : adj.getOrDefault(u, Collections.emptyList())) {
            //自环
            if (Objects.equals(u, v)) {
                circle.add(u);
                circle.add(v);
                break;
            }
            //大环
            else if (visit.getOrDefault(v, VisitState.WHITE) == VisitState.GRAY) {
                circle.add(v);
                T tmp = u;
                while (tmp != null && !Objects.equals(tmp, v)) {
                    circle.add(tmp);
                    tmp = getFather0(tmp);
                }
                if (tmp != null) {
                    addCircle0(tmp);
                }
                break;
            } else if (visit.getOrDefault(v, VisitState.WHITE) == VisitState.WHITE) {
                addFather0(v, u);
                dfs(v);
            }
        }
        visit.put(u, VisitState.BLACK);
    }

    public T getFather(T child) {
        if (child == null) {
            return null;
        }
        return getFather0(child);
    }

    private void addFather0(T child, T father) {
        child2Father.put(child, father);
    }

    private T getFather0(T child) {
        return child2Father.get(child);
    }

    private void addCircle0(T node) {
        circle.add(node);
    }
}
